import numpy as np
import cv2
import os

# TRAIN_DATA_PATH = "../DV_Dataset_All/"
TRAIN_DATA_PATH = "/data2/cyx/DV_Dataset_All_v2/Train/"
# SCALE = 4

def readPatchData(vidIndex, imsize=[200,400], scale=2):
    frame_num = len(os.listdir(TRAIN_DATA_PATH + "data/0/"))
    # imsize = [50, 100]# cv2.imread(TRAIN_DATA_PATH + "/data/0/0.jpg").shape[:2]
    
    x = np.zeros([frame_num, 3, imsize[0]//scale, imsize[1]//scale])
    y = np.zeros([1, 3, imsize[0], imsize[1]])

    fileName = TRAIN_DATA_PATH + "data/" + str(vidIndex) + "/" + str(1) + ".jpg"
    img = cv2.cvtColor(cv2.imread(fileName), cv2.COLOR_BGR2RGB)
    random_row = np.random.randint(img.shape[0] - imsize[0]//2, size = 1)
    random_col = np.random.randint(img.shape[1] - imsize[1]//2, size = 1)

    for i in range(frame_num):
        fileName = TRAIN_DATA_PATH + "data/" + str(vidIndex) + "/" + str(i) + ".jpg"
        img = cv2.cvtColor(cv2.imread(fileName), cv2.COLOR_BGR2RGB)
        # x[i] = cv2.resize(img,(imsize[1]//2, imsize[0]//2)).transpose(2,0,1) / 255.0
        img_crop = img[random_row[0]:random_row[0]+imsize[0]//2,random_col[0]:random_col[0]+imsize[1]//2,:]
        if scale != 2:
            img_crop = cv2.resize(img_crop, (imsize[1]//scale, imsize[0]//scale))
        x[i] = img_crop.transpose(2,0,1) / 255.0

    
    img = cv2.cvtColor(cv2.imread(TRAIN_DATA_PATH + "label/" + str(vidIndex) + ".jpg"), cv2.COLOR_BGR2RGB)
    #y[0] = cv2.resize(img, (imsize[1], imsize[0])).transpose(2,0,1) / 255.0
    y[0] = img[2 * random_row[0]: 2 * random_row[0]+imsize[0],2 * random_col[0]:2 * random_col[0]+imsize[1],:].transpose(2,0,1) / 255.0
    
    return x, y

def generateBatch_FromSavedFiles(batchIndex, BATCH_SIZE=4, imsize=[200,400], scale=2):
    frame_num = len(os.listdir(TRAIN_DATA_PATH + "data/0/"))
    # imsize = [50, 100]# cv2.imread(TRAIN_DATA_PATH + "data/0/0.jpg").shape[:2]
    
    x = np.zeros([BATCH_SIZE, frame_num, 3, imsize[0]//scale, imsize[1]//scale]).astype("float32")
    y = np.zeros([BATCH_SIZE, 1, 3, imsize[0], imsize[1]]).astype("float32")
    for i in range(batchIndex*BATCH_SIZE, (batchIndex+1)*BATCH_SIZE):
        x[i%BATCH_SIZE],  y[i%BATCH_SIZE] = readPatchData(batchIndex, imsize=imsize, scale=scale)
    
    return x, y

if __name__=="__main__":
    x, y = generateBatch_FromSavedFiles(4)
    print(x.shape)